# 工具类
import os
import random
from shutil import copy2


def data_set_split(src_data_folder, src_label_folder, target_data_folder, train_scale=0.7, val_scale=0.2):
    '''
    读取源数据文件夹，生成划分好的文件夹，分为train、val、test三个文件夹进行
    :param src_data_folder: 源文件夹
    :param target_data_folder: 目标文件夹
    :param train_scale: 训练集比例
    :param val_scale: 验证集比例
    :param test_scale: 测试集比例
    :return:
    '''
    print("开始数据集划分")
    img_path_list = os.listdir(src_data_folder)
    # 在目标目录下创建文件夹
    split_names = ['train', 'val', 'test']
    for split_name in split_names:
        split_path = os.path.join(target_data_folder, split_name)
        if os.path.isdir(split_path):
            pass
        else:
            os.makedirs(split_path)
            os.makedirs(os.path.join(split_path, "images"))
            os.makedirs(os.path.join(split_path, "labels"))

    # 按照比例划分数据集，并进行数据图片的复制
    # 首先进行分类遍历

    img_list_length = len(img_path_list)
    train_stop_flag = img_list_length * train_scale
    val_stop_flag = img_list_length * (train_scale + val_scale)
    current_idx = 0
    train_num = 0
    test_num = 0
    val_num = 0

    for index in range(img_list_length):
        current_image_path = os.path.join(src_data_folder, img_path_list[index])
        label_name = img_path_list[index].replace('.jpg', '.txt').replace('.png', '.txt')
        current_label_path = os.path.join(src_label_folder, label_name)

        train_folder = os.path.join(os.path.join(os.path.join(target_data_folder, 'train'), 'images'),
                                    img_path_list[index])
        train_label_folder = os.path.join(os.path.join(os.path.join(target_data_folder, 'train'), 'labels'), label_name)
        val_folder = os.path.join(os.path.join(os.path.join(target_data_folder, 'val'), 'images'), img_path_list[index])
        val_label_folder = os.path.join(os.path.join(os.path.join(target_data_folder, 'val'), 'labels'), label_name)
        test_folder = os.path.join(os.path.join(os.path.join(target_data_folder, 'test'), 'images'),
                                   img_path_list[index])
        test_label_folder = os.path.join(os.path.join(os.path.join(target_data_folder, 'test'), 'labels'), label_name)

        if current_idx <= train_stop_flag:
            copy2(current_image_path, train_folder)
            copy2(current_label_path, train_label_folder)
            train_num = train_num + 1
        elif current_idx <= val_stop_flag:
            copy2(current_image_path, val_folder)
            copy2(current_label_path, val_label_folder)
            val_num = val_num + 1
        else:
            copy2(current_image_path, test_folder)
            copy2(current_label_path, test_label_folder)
            test_num = test_num + 1

        current_idx = current_idx + 1

    print("**********************************************************************")
    print("训练集{}张".format(train_num))
    print("验证集{}张".format(val_num))
    print("测试集{}张".format(test_num))


if __name__ == '__main__':
    src_data_folder = "E:/pycharm/Codes/img/image"
    src_label_folder = "E:/pycharm/Codes/img/label"
    target_data_folder = "E:/pycharm/Codes/img/next_img"
    data_set_split(src_data_folder, src_label_folder, target_data_folder)
